import torch
import copy
from abc import ABC, abstractmethod
from typing import Dict, Any
from tqdm import tqdm


class BaseTuner(ABC):
    def __init__(self, model, train_dataloader, val_dataloader, config: Dict[str, Any]):
        """
        初始化 BaseTuner
        :param model: 需要调整的 Transformer 模型
        :param train_dataloader: 训练数据加载器
        :param val_dataloader: 验证数据加载器
        :param config: 超参数配置
        """
        self.model = model
        self.train_dataloader = train_dataloader
        self.val_dataloader = val_dataloader
        self.config = config
        self.best_model = None
        self.best_score = float('-inf')
        self.optimizer = self._get_optimizer(config.get("optimizer", "AdamW"), config.get("learning_rate", 1e-3))

    def generate_masks(self, src, tgt):
        """生成掩码"""
        src_mask = (src != 0).unsqueeze(1).unsqueeze(2)
        seq_len = tgt.size(1)
        nopeak_mask = torch.tril(torch.ones(seq_len, seq_len, device=tgt.device)).bool()
        tgt_mask = ((tgt != 0).unsqueeze(1).unsqueeze(2)) & nopeak_mask.unsqueeze(0)
        return src_mask, tgt_mask

    @abstractmethod
    def tune(self, optimizer_name: str = "AdamW", epochs: int = 10):
        """
        进行超参数调整
        :param optimizer_name: 优化器名称
        :param epochs: 训练轮数
        """
        pass

    def _get_optimizer(self, optimizer_name: str, learning_rate: float):
        """
        根据优化器名称获取优化器实例
        """
        if optimizer_name == "AdamW":
            return torch.optim.AdamW(self.model.parameters(), lr=learning_rate)
        elif optimizer_name == "SGD":
            return torch.optim.SGD(self.model.parameters(), lr=learning_rate, momentum=0.9)
        else:
            raise ValueError(f"Unsupported optimizer: {optimizer_name}")

    def update_hyperparameters(self, new_config: Dict[str, Any]):
        """
        更新超参数
        :param new_config: 新超参数配置
        """
        self.config.update(new_config)
        self._apply_hyperparameters()

    def _apply_hyperparameters(self):
        """
        应用超参数到模型和优化器
        """
        for param, value in self.config.items():
            if param == "learning_rate":
                for param_group in self.optimizer.param_groups:
                    param_group["lr"] = value
            elif param == "dropout":
                self._set_dropout(self.model, value)

    def _set_dropout(self, module, dropout_value):
        """
        递归设置模型中的 dropout 值
        """
        for child in module.children():
            if isinstance(child, torch.nn.Dropout):
                child.p = dropout_value
            self._set_dropout(child, dropout_value)


